import pymongo
import re
from pymongo import MongoClient

client = pymongo.MongoClient(host='127.0.0.1', port=27017)
# db 表示数据库名称，username 用户名 password 密码
#client.db.authenticate(useranme, password)
# 建立数据库连接
db = client.CNNVD_vul # or db = client['test']
collection = db.vuldb # or collection = db['students']


#传统模糊匹配检索
def search_name_regex(string):
    results = collection.find({"vulName": {"$regex": string, "$options": "i"}})

    # 输出检索结果
    for result in results:
        print(result["vulName"])

#根据漏洞名称
def search_name_cnnvdcode(string):
    #查询，并且返回特定的属性，id默认返回，可以加上"_id":0
    res = collection.find({"vulName":string} , {"vulName":1,"cnnvdCode": 1,"cveCode":1,"vulTypeName":1,"affectedVendor":1, "vulDesc":1, "referUrl":1})
    ldata=[]
    #查询到的结果输出
    for doc in res:
        # 对每个文档执行操作
        vulName = doc["vulName"]
        cnnvdCode = doc["cnnvdCode"]
        cveCode = doc["cveCode"]
        vulTypeName = doc["vulTypeName"]
        affectedVendor = doc["affectedVendor"]
        vulDesc = doc["vulDesc"]
        referUrl = doc["referUrl"]
        
        # 打印文档中的某些字段
        ldata.append(vulName)
        ldata.append(cnnvdCode)
        ldata.append(cveCode)
        ldata.append(vulTypeName)
        ldata.append(affectedVendor)
        ldata.append(vulDesc)
        ldata.append(referUrl)
    return ldata

#根据漏洞cnnvdcode编号查询漏洞的相信信息
def search_vul(string):
    #搜索
    res = collection.find({'cnnvdCode': string})
    #转化为列表输出
    lres=list(res)
    print(lres)




if __name__ == '__main__':
    print("start")
    print("全匹配：")
    search_name_regex("ISC BIND 9 函数dns_db_findrdataset() 拒绝服务漏洞")
    print("去尾匹配：")
    search_name_regex("ISC BIND 9 函数dns_db_findrdataset()")
    print("掐头匹配：")
    search_name_regex("BIND 9 函数dns_db_findrdataset() 拒绝服务漏洞")
    print("掐头去尾匹配：")
    search_name_regex("函数dns_db_findrdataset()")
    print("错误但接近的匹配：")
    search_name_regex("ISC BIND 996 函数dns_db_findrdataset() 拒绝服务漏洞")


    

